# 用streamlit 构建一个问答界面
import streamlit as st
from langchain.memory import ConversationSummaryBufferMemory
from interface.qa_agent import qa_agent
from langchain_community.llms import Tongyi
from settings import settings

def setup_streamlit_interface():
    """
    设置Streamlit界面，包括标题
    """
    # 设置页面标题
    st.title("📑 本地智能客服问答系统")

    # 侧边栏获取API密钥
    with st.sidebar:
        st.write("### 确保在settings.py中设置API密钥")
    api_key = settings.DASH_API_KEY
    model = Tongyi(dashscope_api_key=api_key)

    # 初始化session_state，如果第一次运行
    if "memory" not in st.session_state:
        st.session_state["memory"] = ConversationSummaryBufferMemory(
            llm=model,
            memory_key="chat_history",
            return_messages=True,
            output_key="answer"
        )
    if "messages" not in st.session_state:
        st.session_state["messages"] = [
            {"role": "assistant", "content": "你好，我是AI智能客服，有什么可以帮你的吗？"}
        ]

def display_conversation():
    # 显示历史对话记录
    for message in st.session_state["messages"]:
        st.chat_message(message["role"]).write(message["content"])

def handle_user_input():
    # 处理用户输入
    prompt = st.chat_input("请输入您的问题:")
    if prompt:
        st.session_state["messages"].append({"role": "user", "content": prompt})
        st.chat_message("user").write(prompt)

        with st.spinner("AI正在思考中，请稍等..."):
            response = qa_agent(prompt, settings.DASH_API_KEY, st.session_state["memory"])
        # 确保 response 是一个字典并且包含 'answer' 键
        if isinstance(response, dict) and 'answer' in response:
            answer = response['answer']
            st.session_state["messages"].append({"role": "assistant", "content": answer})
            st.chat_message("assistant").write(answer)
        else:
            raise ValueError("Unexpected response format from qa_agent.")

def main():
    setup_streamlit_interface()
    display_conversation()
    handle_user_input()

# if __name__ == "__main__":
#     main()